from http.server import ThreadingHTTPServer, BaseHTTPRequestHandler
from struct import pack, unpack
from hashlib import md5, sha1
from platform import architecture
from urllib.parse import parse_qs
import datetime, html, os, json, threading, sqlite3, sys, re, getopt


host = {'ip': '0.0.0.0', 'port': 6912}

domain = "http://www.wpfly.cn"

#管理接口的签名密码
secret = "7507f247fec5750cc8dfc69d9eba053b2e085b3e"

#链接序号的加密密码（修改此密码必须重建数据库，否则可能与之前存在的key冲突，因此不建议中途修改）
password = "0553db9fb960e0ae452de31d52d81a9eb90a4ad8"

bit_number = architecture()[0]
if bit_number != "64bit":
    print("The server only supports running on 64 bit systems")

current_path = os.path.dirname(os.path.abspath(__file__))
dbfile = os.path.join(current_path, "shortlink.db")

if not os.path.exists(dbfile):
    os.open(dbfile, os.O_CREAT)
if not os.path.isfile(dbfile):
    print("Cannot found database file: %s" % (dbfile))
    sys.exit()

con = sqlite3.connect(dbfile, check_same_thread=False)
con.row_factory = sqlite3.Row
cur = con.cursor()
cur.execute("CREATE TABLE IF NOT EXISTS link (key TEXT (20) DEFAULT ('') PRIMARY KEY, url TEXT (1000) DEFAULT (''), expires INTEGER DEFAULT (0), type INTEGER DEFAULT (0), addtime INTEGER DEFAULT (0));")
cur.execute("CREATE TABLE IF NOT EXISTS sn (id INTEGER PRIMARY KEY AUTOINCREMENT, stage INTEGER UNIQUE);")
cur.execute("INSERT OR IGNORE INTO link(key, url, expires, type) VALUES(?, ?, ?, ?);", ("0", "https://www.wpfly.cn/", 0, 0))
con.commit()    
    
class RequestHandler(BaseHTTPRequestHandler):

    timeout = 3

    def version_string(self):
        return "SHORTLINK v1.0"
    
    def get_response(self, info):
        self.send_response(200)
        self.send_header("Content-type", "text/html; charset=utf-8")
        self.end_headers()
        html_str = '''<!doctype html>
<html>
<head>
<meta charset="utf-8">
<meta name="viewport" content="width=device-width, initial-scale=1.0">
<title>ShortLink</title>
</head>
<body>
<h1>%s</h1>
</body></html>''' % info
        self.wfile.write(html_str.encode()) 
        
    def do_GET(self):
        
        #请求地址
        path = self.path.rstrip('/')
        match_obj = re.match(r"^/([A-Za-z0-9]{1,20})$", path)
        if match_obj is None :
            return self.get_response(("链接格式错误！"))
        res = cur.execute("select url, expires from link where key=?", (match_obj.groups()[0],))  
        row =  res.fetchone()
        if row is None :
            return self.get_response(("链接不存在！请检查输入，注意区分大小写。"))
        #expires=0：永久有效
        expires = int(row["expires"])
        if expires != 0 and expires < int(datetime.datetime.now().timestamp()) :
            return self.get_response(("链接已失效！"))
        self.send_response(302)
        self.send_header("Location", row["url"])
        self.end_headers()        

    def post_success(self, data = None):
        self.wfile.write(json.dumps(dict(code = 0, msg = "success", data = data), ensure_ascii = False).encode())
    
    def post_failure(self, code = 1, msg = None, data = None):
        code = int(code)
        if code < 1:
            code = 1
        errorMsg = {
            "ERROR_1" : "系统异常，请稍后再试",  
            "ERROR_10001" : "请求地址错误",
            "ERROR_10002" : "参数缺失",
            "ERROR_10003" : "参数格式错误",
            "ERROR_10004" : "url无效",
            "ERROR_10005" : "expires无效",
            "ERROR_10006" : "",
            "ERROR_10007" : "签名验证失败",
            "ERROR_10008" : "签名过期",
            "ERROR_10009" : "数据保存失败",
        }
        if msg is None:
            key = "ERROR_" + str(code)
            msg = errorMsg[key] if key in errorMsg else "未定义的错误码"
        self.wfile.write(json.dumps(dict(code = code, msg = msg, data = data), ensure_ascii = False).encode())           
        
    def verify_signature(self, data, secret, field):
        keys =  list(data.keys())
        if field not in keys:
            return False
        keys.remove(field)
        keys.sort()
        kv = list()
        for k in keys:
            kv.append("%s=%s" % (k, data[k][0]))
        #签名密钥附加到数据末尾
        kv.append(secret)
        signature = sha1('&'.join(kv).encode()).hexdigest()
        #print(signature)
        return True if str(data[field][0]).lower() == signature else False

    def to62(self, num):
        base62_charset = '0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ';
        ret = ''
        while(num >= 0):
            c = num % 62
            ret += base62_charset[c]
            num = num // 62
            if num == 0 :
                break
        return ret

    def from62(self, num):
        base62_charset = '0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ';
        num_len = len(num)
        ret = 0
        for i in range(num_len):
            pos = base62_charset.index(num[i])
            ret += (pos * (62 ** i))
        return ret
    
    def encrypt(self, data, password):
        key_bytes = md5(password.encode()).digest()
        data_bytes = bytearray(pack('>Q', data))
        data_len = len(data_bytes)
        for i in range((2*data_len) -1):
            src = i % data_len 
            dst = (i + 1) % data_len
            k = i & 15
            data_bytes[dst] = ((data_bytes[src] + data_bytes[dst]) & 255) ^ key_bytes[k]
        return self.to62(unpack('>Q', data_bytes)[0])

    def decrypt(self, data, password):
        key_bytes = md5(password.encode()).digest()
        data_bytes = bytearray(pack('>Q', self.from62(data)))
        data_len = len(data_bytes)
        for i in range((2*data_len) - 2, -1 , -1):
            src = i % data_len 
            dst = (i + 1) % data_len
            k = i & 15
            data_bytes[dst] = ((data_bytes[dst] ^ key_bytes[k]) + 256 - data_bytes[src]) & 255
        return unpack('>Q', data_bytes)[0]
    
    def do_POST(self):        
        
        self.send_response(200)
        self.send_header("Content-type", "application/json; charset=utf-8")
        self.end_headers()
        
        #请求地址
        path = self.path.rstrip('/')
        allow_request = ["/api/link"]
        if path not in allow_request:
            return self.post_failure(10001)
    
        #获取POST提交的数据
        datas = self.rfile.read(int(self.headers['content-length'])).decode()
        datas = parse_qs(datas, True)
        
        #参数检查（正则）
        need_params = {"url": r"^\S{1,1000}$", "expires": r"^\d{1,10}$", "type": r"^\d$", "timestamp": r"^\d{1,10}$", "signature": r"^[0-9a-fA-F]{1,40}$"}
        for k, v in need_params.items():
            if k not in datas:
                return self.post_failure(10002, None, k)            
            if re.match(v, datas[k][0]) is None:
                return self.post_failure(10003, None, k)
        
        p_url = datas["url"][0]
        p_expires = int(datas["expires"][0])
        p_type = int(datas["type"][0])
        p_timestamp = int(datas["timestamp"][0])
        
        if not p_url.startswith("http://") and not p_url.startswith("https://"):
            return self.post_failure(10004)
        if p_expires != 0 and p_expires < int(datetime.datetime.now().timestamp()):
            return self.post_failure(10005)        
        
        #签名验证
        if not self.verify_signature(datas, secret, 'signature'):
            return self.post_failure(10007)
        
        #签名时效验证：3分钟       
        if p_timestamp + 180 < int(datetime.datetime.now().timestamp()):
            return self.post_failure(10008)
            
        #保存数据并更新Key
        lock = threading.Lock()
        lock.acquire()
        linkId = ""
        try:
            cur.execute("Replace into sn(stage) values(?)", (0,))
            con.commit()
            if not cur.lastrowid:
                raise Exception("链接序号生成失败")
            linkId = self.encrypt(cur.lastrowid, password)  
            cur.execute("Insert into link(key, url, expires, type, addtime) values(?,?,?,?,?)", (linkId, p_url, p_expires, p_type, int(datetime.datetime.now().timestamp())))
            con.commit()            
        except Exception as e:
            #print(e)
            lock.release()
            return self.post_failure(10009)
        
        #正常释放锁
        lock.release() 
        
        #返回url, domain及path
        return self.post_success({
            "url": domain + linkId,
            "domain": domain,
            "path": linkId,
        })           
            
        self.wfile.write(self.post_failure().encode())

    
if __name__ == '__main__':
    
    try:
        opts, args = getopt.getopt(sys.argv[1:], "hi:p:d:s:", ["help", "ip=", "port=", "domain=", "secret="])
    except getopt.GetoptError:
        print('python3 shortlink.py -h')
        sys.exit(2)
        
    for opt, arg in opts:
        if opt in ("-h", "--help"):
            print("[usage]:")
            print("python3 shortlink.py -h")
            print("python3 shortlink.py -i [ip] -p [port] -d [domain] -s [secret]")    
            print("python3 shortlink.py --ip=[ip] --port=[port] --domain=[domain] --secret=[secret]")
            sys.exit()
        elif opt in ("-i", "--ip"):
            host["ip"] = arg
        elif opt in ("-p", "--port"):
            host["port"] = arg
        elif opt in ("-d", "--domain"):
            domain = arg
        elif opt in ("-s", "--secret"):
            secret = arg
    
    if re.match(r"^((\d)|([1-9]\d)|(1\d{2})|(2[0-4]\d)|(25[0-5]))\.((\d)|([1-9]\d)|(1\d{2})|(2[0-4]\d)|(25[0-5]))\.((\d)|([1-9]\d)|(1\d{2})|(2[0-4]\d)|(25[0-5]))\.((\d)|([1-9]\d)|(1\d{2})|(2[0-4]\d)|(25[0-5]))$", host["ip"]) is None:
        host["ip"] = '0.0.0.0'
    host["port"] = int(host["port"])
    if len(domain) > 0 and not domain.startswith("http://"):
        domain = "http://" + domain
    domain = domain.rstrip('/') + "/"
    
    httpd = ThreadingHTTPServer((host["ip"], host["port"]), RequestHandler)
    print("Starting server, listen at: %s:%s" % (host["ip"], host["port"]))
    httpd.serve_forever()
